from fastapi import WebSocket
from typing import List, Dict


class WebsocketManager:
    def __init__(self):
        self.active_connection = None

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connection = websocket

    def disconnect(self):
        self.active_connection = None

    async def send_text(self, message: str):
        await self.active_connection.send_text(message)

    async def send_json(self, message: dict):
        await self.active_connection.send_json(message)


# 多个客户端请求
class ConnectionManager:
    def __init__(self):
        # 存放活跃链接
        self.active_connections: List[Dict[str, WebSocket]] = []

    async def connect(self, client: str, ws: WebSocket):
        # 客户端连接服务器端
        await ws.accept()
        self.active_connections.append({"client": client, "ws": ws})

    def disconnect(self, client: str, ws: WebSocket):
        # 断开连接，移除 ws 对象
        self.active_connections.remove({"client": client, "ws": ws})

    async def send_json(self, client: str, message: dict):
        # 发送 json 格式数据到特定客户端
        for connection in self.active_connections:
            if connection["client"] == client:
                await connection["ws"].send_json(message)

    async def send_txt(self, client: str, message: dict):
        # 发送 json 格式数据到特定客户端
        for connection in self.active_connections:
            if connection["client"] == client:
                await connection["ws"].send_text(message)

    async def broadcast_json(self, message: dict):
        # 广播消息到所有客户端
        for connection in self.active_connections:
            await connection["ws"].send_json(message)